from typing import Dict, List, Optional, Set
from datetime import datetime
from loguru import logger

class OSVSearcher:
    def __init__(self, vulnerabilities: List[Dict]):
        self.vulnerabilities = vulnerabilities
        self._build_indexes()
        
    def _build_indexes(self):
        """构建索引以提高搜索性能"""
        self._cwe_index = {}
        self._cve_index = {}
        self._package_index = {}
        self._severity_index = {}
        self._source_index = {}
        
        for vuln in self.vulnerabilities:
            # CWE索引
            for cwe in vuln['cwes']:
                if cwe not in self._cwe_index:
                    self._cwe_index[cwe] = []
                self._cwe_index[cwe].append(vuln)
            
            # CVE索引
            for alias in vuln['aliases']:
                if alias.startswith('CVE-'):
                    if alias not in self._cve_index:
                        self._cve_index[alias] = []
                    self._cve_index[alias].append(vuln)
            
            # 包索引
            for package in vuln['packages']:
                key = (package['ecosystem'], package['name'])
                if key not in self._package_index:
                    self._package_index[key] = []
                self._package_index[key].append(vuln)
            
            # 严重程度索引
            for severity in vuln['severity']:
                sev_type = severity.get('type')
                if sev_type:
                    if sev_type not in self._severity_index:
                        self._severity_index[sev_type] = []
                    self._severity_index[sev_type].append(vuln)
            
            # 来源索引
            source = vuln.get('source', 'unknown')
            if source not in self._source_index:
                self._source_index[source] = []
            self._source_index[source].append(vuln)
    
    def search_by_cwe(self, cwe_id: str) -> List[Dict]:
        """按CWE ID搜索漏洞"""
        return self._cwe_index.get(cwe_id, [])
    
    def search_by_cve(self, cve_id: str) -> List[Dict]:
        """按CVE ID搜索漏洞"""
        return self._cve_index.get(cve_id, [])
    
    def search_by_package(self, ecosystem: str, package_name: str) -> List[Dict]:
        """按包名搜索漏洞"""
        return self._package_index.get((ecosystem, package_name), [])
    
    def search_by_severity(self, severity: str) -> List[Dict]:
        """按严重程度搜索漏洞"""
        return self._severity_index.get(severity, [])
    
    def search_by_source(self, source: str) -> List[Dict]:
        """按数据来源搜索漏洞"""
        return self._source_index.get(source, [])
    
    def search_by_date_range(self, start_date: datetime, end_date: datetime) -> List[Dict]:
        """按日期范围搜索漏洞"""
        results = []
        for vuln in self.vulnerabilities:
            if vuln['published']:
                published_date = datetime.fromisoformat(vuln['published'].replace('Z', '+00:00'))
                if start_date <= published_date <= end_date:
                    results.append(vuln)
        return results
    
    def search_by_keyword(self, keyword: str) -> List[Dict]:
        """按关键词搜索漏洞"""
        results = []
        keyword = keyword.lower()
        for vuln in self.vulnerabilities:
            # 搜索摘要
            if keyword in vuln['summary'].lower():
                results.append(vuln)
                continue
                
            # 搜索详情
            if keyword in vuln['details'].lower():
                results.append(vuln)
                continue
                
            # 搜索包名
            for package in vuln['packages']:
                if keyword in package['name'].lower():
                    results.append(vuln)
                    break
        return results
    
    def advanced_search(self, 
                       cwe_id: Optional[str] = None,
                       cve_id: Optional[str] = None,
                       ecosystem: Optional[str] = None,
                       package_name: Optional[str] = None,
                       severity: Optional[str] = None,
                       source: Optional[str] = None,
                       start_date: Optional[datetime] = None,
                       end_date: Optional[datetime] = None,
                       keyword: Optional[str] = None) -> List[Dict]:
        """高级搜索"""
        # 使用索引进行快速过滤
        results = set()
        
        if cwe_id:
            results.update(self._cwe_index.get(cwe_id, []))
        elif cve_id:
            results.update(self._cve_index.get(cve_id, []))
        elif ecosystem and package_name:
            results.update(self._package_index.get((ecosystem, package_name), []))
        elif severity:
            results.update(self._severity_index.get(severity, []))
        elif source:
            results.update(self._source_index.get(source, []))
        else:
            results = set(self.vulnerabilities)
        
        # 转换为列表以便进一步处理
        results = list(results)
        
        # 应用额外的过滤条件
        if start_date and end_date:
            results = [v for v in results if v['published'] and 
                      start_date <= datetime.fromisoformat(v['published'].replace('Z', '+00:00')) <= end_date]
            
        if keyword:
            keyword = keyword.lower()
            results = [v for v in results if 
                      keyword in v['summary'].lower() or
                      keyword in v['details'].lower() or
                      any(keyword in p['name'].lower() for p in v['packages'])]
        
        return results
    
    def get_related_vulnerabilities(self, vuln_id: str) -> List[Dict]:
        """获取相关漏洞"""
        target_vuln = None
        for vuln in self.vulnerabilities:
            if vuln['id'] == vuln_id:
                target_vuln = vuln
                break
                
        if not target_vuln:
            return []
            
        related = set()
        
        # 通过CWE关联
        for cwe in target_vuln['cwes']:
            related.update(self._cwe_index.get(cwe, []))
                
        # 通过包名关联
        for package in target_vuln['packages']:
            key = (package['ecosystem'], package['name'])
            related.update(self._package_index.get(key, []))
        
        # 移除目标漏洞本身
        related.discard(target_vuln)
        
        return list(related)
    
    def get_unique_values(self, field: str) -> Set[str]:
        """获取指定字段的唯一值"""
        values = set()
        for vuln in self.vulnerabilities:
            if field in vuln:
                if isinstance(vuln[field], list):
                    values.update(vuln[field])
                else:
                    values.add(vuln[field])
        return values 